import os
import sys
from datetime import datetime, timedelta, timezone

import torch
import torch.distributed as dist


import utils.env
from utils.overwatch import initialize_overwatch


from transformers import (
    HfArgumentParser,
    TrainingArguments,
    logging as hf_logging
)
import datasets

from quantize.config import OtherArguments, QuantArguments,QATArguments







def parse_args(
    print_args=False,
) -> tuple[TrainingArguments, QATArguments, QuantArguments, OtherArguments]:
    parser = HfArgumentParser(
        (
            TrainingArguments,
            QATArguments,
            QuantArguments,
            OtherArguments,
        )  # type: ignore
    )
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        training_args, qat_args, quant_args, other_args = parser.parse_json_file(
            json_file=os.path.abspath(sys.argv[1])
        )
    else:
        training_args, qat_args, quant_args, other_args = (
            parser.parse_args_into_dataclasses()
        )
    if print_args:
        print(
            "Config:",
            *(f"{i}\n" for i in (training_args, qat_args, quant_args, other_args)),
        )
    return training_args, qat_args, quant_args, other_args


def main(state):
    utils.env.set_seed(42)
    logger = initialize_overwatch("vla_qat")
    training_args, qat_args, quant_args, other_args = parse_args()
    if not other_args.test and not any(
        (
            other_args.eval_quant_fp,
            other_args.eval_quant_int,
            other_args.save_int_tensors,
        )
    ):
        if qat_args.resume_dir is not None and qat_args.resume_dir.exists():
            task_name = qat_args.resume_dir.parent.name
            run_dir = qat_args.resume_dir.parent
        else:
            now = datetime.now(timezone.utc) + timedelta(hours=8)
            task_name = (
                f"vla_q+{qat_args.dataset_name}"
                f"-{now.strftime(r'%y%m%d_%H:%M')}"
                f"+b{training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps}"
                f"+lr-{training_args.learning_rate}"
            )

            run_dir = qat_args.run_root_dir / task_name
            run_dir.mkdir(exist_ok=True, parents=True)
        training_args.output_dir = str(run_dir.absolute())

        logger.info(f"Run dir => {training_args.output_dir}")
    else:
        task_name = "test"
        run_dir = None

    # def on_exit(sig=None, frame=None):
    #     print(f"\033[33mRun dir => {training_args.output_dir}\033[0m")
    #     # exit(0)
    # original_sigint_handler = signal.getsignal(signal.SIGINT)
    # original_sigterm_handler = signal.getsignal(signal.SIGTERM)
    # signal.signal(
    #     signal.SIGINT, lambda *args, **kwargs: on_exit() or original_sigint_handler()
    # )  # Ctrl+C
    # signal.signal(
    #     signal.SIGTERM, lambda *args, **kwargs: on_exit() or original_sigterm_handler()
    # )  # 终止信号
    # atexit.register(on_exit)
    
    match qat_args.mode:
        case "qat":
            from qat import qat_main
            qat_main(
                task_name, run_dir, training_args, qat_args, quant_args, other_args, state=state
            )
        case "spin":
            from test_spin_quant import test_main
            test_main(
                task_name, run_dir, training_args, qat_args, quant_args, other_args, state=state
            )
        




if __name__ == "__main__":
    try:
        main(state=None)
    finally:
        dist.destroy_process_group()
